import os
import wget
import zipfile
import tarfile
from tqdm import tqdm

# 数据集下载配置
DATASETS = {
    "DIV2K": {
        "train": "http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_HR.zip",
        "valid": "http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_valid_HR.zip"
    },
    "Flickr2K": {
        "url": "https://cv.snu.ac.kr/research/EDSR/Flickr2K.tar"
    }
}

def download_dataset(url, save_dir, filename=None):
    """带进度条的数据集下载函数"""
    # 根据输入参数save_dir，新建存储目录
    os.makedirs(save_dir, exist_ok=True)
    
    # 传入文件名参数检查
    if not filename:
        filename = url.split("/")[-1]
    save_path = os.path.join(save_dir, filename)
    
    # 文件存在性检查
    if os.path.exists(save_path):
        print(f"文件 {filename} 已存在，跳过下载")
        return save_path

    print(f"开始下载 {filename}...")
    """
    @date 2025/09/29
    @author cmx-cxd
    移除下载进度的无效重复行输出（会导致日志行数过多，文件过大，大部分都是无效日志），修正为单行进度条，减少输出频率
    """
    # 使用tqdm创建进度条，减少输出频率
    with tqdm(total=100, desc=f"下载 {filename}", unit="%") as pbar:
        def progress_bar(current, total, width=80):
            progress = current / total * 100
            pbar.n = int(progress)
            pbar.update(0)      # 只更新进度条显示，不产生新输出
            
        wget.download(url, save_path, bar=progress_bar)
    # def progress_bar(current, total, width = 80):
    #     progress = current / total * 100
    #     # 打印下载进度条
    #     tqdm.write(f"下载进度：{progress:.2f}%", end="\r")
        
    # 不再使用自定义进度条执行下载
    # wget.download(url, save_path, bar=progress_bar)
    print(f"\n{filename} 下载完成")
    return save_path

def extract_archive(file_path, extract_dir):
    """解压数据集文件 - 添加重复解压检查"""
    # 根据输入参数extract_dir，新建解压目录
    os.makedirs(extract_dir, exist_ok=True)
    """
    @date 2025/09/29
    @author cmx-cxd
    新增重复解压检查，避免二次解压数据集
    1、若压缩包已下载 → 跳过下载（已有检查）。
    2、若已解压且成功 → 跳过解压（新增标记检查）。
    3、若解压中断 / 失败 → 重新解压（标记文件不存在）。
    """
    filename = os.path.basename(file_path)
    # 生成解压后根目录的标识（根据压缩包名判断）
    extract_flag = os.path.join(extract_dir, f".{filename}.extracted")  # 标记文件
    
    # 检查是否已解压（通过标记文件判断）
    if os.path.exists(extract_flag):
        print(f"文件 {filename} 已解压，跳过解压！")
        return
    
    # 执行解压
    try:
        # 文件名后缀检查
        if file_path.endswith(".zip"):
            """处理zip压缩文件"""
            with zipfile.ZipFile(file_path, 'r') as zip_ref:
                # 显示解压进度
                for file in tqdm(zip_ref.namelist(), desc="解压中"):
                    zip_ref.extract(file, extract_dir)
        elif file_path.endswith(".tar") or file_path.endswith(".tar.gz"):
            """处理tar压缩文件"""
            with tarfile.open(file_path, 'r') as tar_ref:
                # 显示解压进度
                members = tar_ref.getmembers()
                for member in tqdm(members, desc="解压中"):
                    tar_ref.extract(member, extract_dir)
                    
        # 解压成功后创建标记文件
        with open(extract_flag, 'w') as f:
            f.write("Extracted successfully")
        print(f"文件 {os.path.basename(file_path)} 解压完成")
    except Exception as e:
        print(f"解压失败：{e}")
        # 失败时删除标记文件（避免误判）
        if os.path.exists(extract_flag):
            os.remove(extract_flag)
        
def prepare_large_datasets(base_dir):
    """准备所有大型数据集"""
    # 下载DIV2K
    div2k_dir = os.path.join(base_dir, "DIV2K")
    for split, url in DATASETS["DIV2K"].items():
        file_path = download_dataset(url, div2k_dir)
        extract_archive(file_path, os.path.join(div2k_dir, split))
        
    # 下载Flickr2K
    flickr_dir = os.path.join(base_dir, "Flickr2K")
    flickr_url = DATASETS["Flickr2K"]["url"]
    file_path = download_dataset(flickr_url, flickr_dir)
    extract_archive(file_path, flickr_dir)
    
    print("所有数据集准备完成")